// Uses nacl's secret_box to encrypt a net.Conn.
// It is (meant to be) an implementation of the STS protocol.
// Note we do not (yet) assume that a remote peer's pubkey
// is known ahead of time, and thus we are technically
// still vulnerable to MITM. (TODO!)
// See docs/sts-final.pdf for more info
package conn

import (
	"bytes"
	crand "crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"io"
	"net"
	"time"

	"golang.org/x/crypto/nacl/box"
	"golang.org/x/crypto/nacl/secretbox"
	"golang.org/x/crypto/ripemd160"

	"github.com/golang/snappy"
	cmn "github.com/lianxiangcloud/linkchain/libs/common"
	"github.com/lianxiangcloud/linkchain/libs/crypto"
	"github.com/lianxiangcloud/linkchain/libs/ser"
)

// 1 + 4 + 32 * 1024 == 32773 total frame size
const (
	leadingSize      = 1
	dataLenSize      = 4
	headerSize       = leadingSize + dataLenSize
	dataMaxSize      = 32 * 1024
	unsealedDataSize = dataMaxSize + dataLenSize
	sealedDataSize   = unsealedDataSize + secretbox.Overhead
	frameCapacity    = 65535

	versionMusk  = 0xF0
	version00    = 0xF0 // version01 = 0xE0, version02 = 0xD0, etc.
	typeMusk     = 0x0F
	typeCompress = 0x0F
	typeEncrypt  = 0x0E // typeX1 = 0xD0, typeX2 = 0xC0, etc.
	typeRaw      = 0x00
)

var (
	leadingVersion byte = version00
	leadingType    byte = typeCompress
)

// Implements net.Conn
type SecretConnection struct {
	conn       io.ReadWriteCloser
	recvBuffer []byte
	recvNonce  *[24]byte //recvNonce from mine is equel with remote sendNonce
	sendNonce  *[24]byte
	remPubKey  crypto.PubKey
	shrSecret  *[32]byte // shared secret
}

// Performs handshake and returns a new authenticated SecretConnection.
// Returns nil if error in handshake.
// Caller should call conn.Close()
// See docs/sts-final.pdf for more information.
func MakeSecretConnection(conn io.ReadWriteCloser, locPrivKey crypto.PrivKey) (*SecretConnection, error) {

	locPubKey := locPrivKey.PubKey()

	// Generate ephemeral keys for perfect forward secrecy.
	locEphPub, locEphPriv := genEphKeys() //generate tmp privatekey and pubeky

	// Write local ephemeral pubkey and receive one too.
	// NOTE: every 32-byte string is accepted as a Curve25519 public key
	// (see DJB's Curve25519 paper: http://cr.yp.to/ecdh/curve25519-20060209.pdf)
	remEphPub, err := shareEphPubKey(conn, locEphPub)
	if err != nil {
		return nil, err
	}

	// Compute common shared secret.
	shrSecret := computeSharedSecret(remEphPub, locEphPriv)

	// Sort by lexical order.
	loEphPub, hiEphPub := sort32(locEphPub, remEphPub)

	// Check if the local ephemeral public key
	// was the least, lexicographically sorted.
	locIsLeast := bytes.Equal(locEphPub[:], loEphPub[:])

	// Generate nonces to use for secretbox.
	recvNonce, sendNonce := genNonces(loEphPub, hiEphPub, locIsLeast)

	// Generate common challenge to sign.
	challenge := genChallenge(loEphPub, hiEphPub)

	// Construct SecretConnection.
	sc := &SecretConnection{
		conn:       conn,
		recvBuffer: nil,
		recvNonce:  recvNonce,
		sendNonce:  sendNonce,
		shrSecret:  shrSecret,
	}

	// Sign the challenge bytes for authentication.
	locSignature := signChallenge(challenge, locPrivKey)

	// Share (in secret) each other's pubkey & challenge signature
	authSigMsg, err := shareAuthSignature(sc, locPubKey, locSignature)
	if err != nil {
		return nil, err
	}
	remPubKey, remSignature := authSigMsg.Key, authSigMsg.Sig
	if remPubKey == nil {
		return nil, errors.New("Peer sent a nil public key")
	}
	
	if !remPubKey.VerifyBytes(challenge[:], remSignature) {
		return nil, errors.New("Challenge verification failed")
	}

	// We've authorized.
	sc.remPubKey = remPubKey
	return sc, nil
}

// Returns authenticated remote pubkey
func (sc *SecretConnection) RemotePubKey() crypto.PubKey {
	return sc.remPubKey
}

// Writes encrypted frames of `sealedFrameSize`
// CONTRACT: data smaller than dataMaxSize is read atomically.
func (sc *SecretConnection) Write(data []byte) (n int, err error) {
	for 0 < len(data) {
		var chunk []byte
		if dataMaxSize < len(data) {
			chunk = data[:dataMaxSize]
			data = data[dataMaxSize:]
		} else {
			chunk = data
			data = nil
		}

		chunkLength := len(chunk)
		var frame []byte = nil

		if leadingType == typeCompress {
			// compress the frame
			frame = make([]byte, headerSize+snappy.MaxEncodedLen(chunkLength))
			frame[0] = leadingVersion | leadingType
			encodedLength := len(snappy.Encode(frame[headerSize:], chunk))
			binary.BigEndian.PutUint32(frame[leadingSize:], uint32(encodedLength))
			frame = frame[:headerSize+encodedLength]

		} else if leadingType == typeEncrypt {
			unsealedData := make([]byte, unsealedDataSize)
			binary.BigEndian.PutUint32(unsealedData, uint32(chunkLength))
			copy(unsealedData[dataLenSize:], chunk)
			// encrypt the frame
			frame = make([]byte, headerSize+sealedDataSize)
			frame[0] = leadingVersion | leadingType
			sealedData := frame[headerSize:]
			secretbox.Seal(sealedData[:0], unsealedData, sc.sendNonce, sc.shrSecret)
			//fmt.Printf("secretbox.Seal(sealed:%v,sendNonce:%X,shrSecret:%X\n", sealedData[:256], sc.sendNonce, sc.shrSecret)
			binary.BigEndian.PutUint32(frame[leadingSize:], uint32(len(sealedData)))
			frame = frame[:headerSize+sealedDataSize]
			incr2Nonce(sc.sendNonce)
			// end encryption

		} else if leadingType == typeRaw {
			// original flow
			unsealedData := make([]byte, unsealedDataSize)
			binary.BigEndian.PutUint32(unsealedData, uint32(chunkLength))
			copy(unsealedData[dataLenSize:], chunk)
			// encrypt the frame
			frame = make([]byte, sealedDataSize)
			secretbox.Seal(frame[:0], unsealedData, sc.sendNonce, sc.shrSecret)
			// fmt.Printf("secretbox.Seal(sealed:%X,sendNonce:%X,shrSecret:%X\n", frame, sc.sendNonce, sc.shrSecret)
			incr2Nonce(sc.sendNonce)
			// end encryption

		} else {
			return 0, errors.New("unknow type to write")
		}

		if len(frame) > frameCapacity {
			return 0, errors.New("write packet length is greater than frameCapacity")
		}

		_, err := sc.conn.Write(frame)
		if err != nil {
			return n, err
		}
		n += len(chunk)
	}
	return
}

// CONTRACT: data smaller than dataMaxSize is read atomically.
func (sc *SecretConnection) Read(data []byte) (n int, err error) {
	if 0 < len(sc.recvBuffer) {
		n = copy(data, sc.recvBuffer)
		sc.recvBuffer = sc.recvBuffer[n:]
		return
	}

	if leadingType == typeRaw {
		// original flow
		return sc.rawRead(data)
	}

	header := make([]byte, headerSize)
	_, err = io.ReadFull(sc.conn, header)
	if err != nil {
		return
	}

	headerVersion := header[0] & versionMusk
	headerType := header[0] & typeMusk
	if headerVersion != version00 || (headerType != typeEncrypt && headerType != typeCompress) {
		return 0, errors.New("Read unknow type packet")
	}

	rawLength := binary.BigEndian.Uint32(header[leadingSize:])
	if rawLength > frameCapacity-headerSize {
		return 0, errors.New("read packet length is greater than frameCapacity")
	}
	rawData := make([]byte, rawLength)
	_, err = io.ReadFull(sc.conn, rawData)
	if err != nil {
		return 0, errors.New("Failed to ReadFull from SecretConnection")
	}

	var frame []byte = nil
	var chunkLength uint32 = 0

	if headerType == typeCompress {
		frame, err = snappy.Decode(nil, rawData)
		if err != nil {
			return 0, errors.New("Failed to decode SecretConnection")
		}
		chunkLength = uint32(len(frame))

	} else if headerType == typeEncrypt {
		// decrypt the frame
		frame = make([]byte, unsealedDataSize)
		//fmt.Printf("secretbox.Open(sealed:%v,recvNonce:%X,shrSecret:%X\n", rawData[:256], sc.recvNonce, sc.shrSecret)
		_, ok := secretbox.Open(frame[:0], rawData, sc.recvNonce, sc.shrSecret)
		if !ok {
			return 0, errors.New("Failed to decrypt SecretConnection")
		}
		incr2Nonce(sc.recvNonce)
		// end decryption
		chunkLength = binary.BigEndian.Uint32(frame)
		frame = frame[dataLenSize:]

	} else {
		// never comes to here
		return 0, errors.New("Read unknow type packet")
	}

	if chunkLength > dataMaxSize {
		return 0, errors.New("chunkLength is greater than dataMaxSize")
	}
	chunk := frame[:chunkLength]

	n = copy(data, chunk)
	sc.recvBuffer = chunk[n:]
	return
}

func (sc *SecretConnection) rawRead(data []byte) (n int, err error) {
	sealedFrame := make([]byte, sealedDataSize)
	_, err = io.ReadFull(sc.conn, sealedFrame)
	if err != nil {
		return
	}

	// decrypt the frame
	var frame = make([]byte, unsealedDataSize)
	// fmt.Printf("secretbox.Open(sealed:%X,recvNonce:%X,shrSecret:%X\n", frame, sc.recvNonce, sc.shrSecret)
	_, ok := secretbox.Open(frame[:0], sealedFrame, sc.recvNonce, sc.shrSecret)
	if !ok {
		return n, errors.New("Failed to decrypt SecretConnection")
	}
	incr2Nonce(sc.recvNonce)
	// end decryption

	var chunkLength = binary.BigEndian.Uint32(frame) // read the first four bytes
	if chunkLength > dataMaxSize {
		return 0, errors.New("chunkLength is greater than dataMaxSize")
	}
	var chunk = frame[dataLenSize : dataLenSize+chunkLength]

	n = copy(data, chunk)
	sc.recvBuffer = chunk[n:]
	return
}

// Implements net.Conn
func (sc *SecretConnection) Close() error                  { return sc.conn.Close() }
func (sc *SecretConnection) LocalAddr() net.Addr           { return sc.conn.(net.Conn).LocalAddr() }
func (sc *SecretConnection) RemoteAddr() net.Addr          { return sc.conn.(net.Conn).RemoteAddr() }
func (sc *SecretConnection) SetDeadline(t time.Time) error { return sc.conn.(net.Conn).SetDeadline(t) }
func (sc *SecretConnection) SetReadDeadline(t time.Time) error {
	return sc.conn.(net.Conn).SetReadDeadline(t)
}
func (sc *SecretConnection) SetWriteDeadline(t time.Time) error {
	return sc.conn.(net.Conn).SetWriteDeadline(t)
}

func genEphKeys() (ephPub, ephPriv *[32]byte) {
	var err error
	ephPub, ephPriv, err = box.GenerateKey(crand.Reader)
	if err != nil {
		panic("Could not generate ephemeral keypairs")
	}
	return
}

func shareEphPubKey(conn io.ReadWriteCloser, locEphPub *[32]byte) (remEphPub *[32]byte, err error) {

	// Send our pubkey and receive theirs in tandem.
	var trs, _ = cmn.Parallel(
		func(_ int) (val interface{}, err error, abort bool) {
			var _, err1 = ser.EncodeWriterWithType(conn, locEphPub)
			if err1 != nil {
				return nil, err1, true // abort
			} else {
				return nil, nil, false
			}
		},
		func(_ int) (val interface{}, err error, abort bool) {
			var _remEphPub [32]byte
			var _, err2 = ser.DecodeReaderWithType(conn, &_remEphPub, 1024*1024) // TODO
			if err2 != nil {
				return nil, err2, true // abort
			} else {
				return _remEphPub, nil, false
			}
		},
	)

	// If error:
	if trs.FirstError() != nil {
		err = trs.FirstError()
		return
	}

	// Otherwise:
	var _remEphPub = trs.FirstValue().([32]byte)
	return &_remEphPub, nil
}

func computeSharedSecret(remPubKey, locPrivKey *[32]byte) (shrSecret *[32]byte) {
	shrSecret = new([32]byte)
	box.Precompute(shrSecret, remPubKey, locPrivKey)
	return
}

func sort32(foo, bar *[32]byte) (lo, hi *[32]byte) {
	if bytes.Compare(foo[:], bar[:]) < 0 {
		lo = foo
		hi = bar
	} else {
		lo = bar
		hi = foo
	}
	return
}

func genNonces(loPubKey, hiPubKey *[32]byte, locIsLo bool) (recvNonce, sendNonce *[24]byte) {
	nonce1 := hash24(append(loPubKey[:], hiPubKey[:]...))
	nonce2 := new([24]byte)
	copy(nonce2[:], nonce1[:])
	nonce2[len(nonce2)-1] ^= 0x01
	if locIsLo {
		recvNonce = nonce1
		sendNonce = nonce2
	} else {
		recvNonce = nonce2
		sendNonce = nonce1
	}
	return
}

func genChallenge(loPubKey, hiPubKey *[32]byte) (challenge *[32]byte) {
	return hash32(append(loPubKey[:], hiPubKey[:]...))
}

func signChallenge(challenge *[32]byte, locPrivKey crypto.PrivKey) (signature crypto.Signature) {
	signature, err := locPrivKey.Sign(challenge[:])
	// TODO(ismail): let signChallenge return an error instead
	if err != nil {
		panic(err)
	}
	return
}

type authSigMessage struct {
	Key crypto.PubKey
	Sig crypto.Signature
}

func shareAuthSignature(sc *SecretConnection, pubKey crypto.PubKey, signature crypto.Signature) (recvMsg authSigMessage, err error) {

	// Send our info and receive theirs in tandem.
	var trs, _ = cmn.Parallel(
		func(_ int) (val interface{}, err error, abort bool) {
			var _, err1 = ser.EncodeWriterWithType(sc, authSigMessage{pubKey, signature})
			if err1 != nil {
				return nil, err1, true // abort
			} else {
				return nil, nil, false
			}
		},
		func(_ int) (val interface{}, err error, abort bool) {
			var _recvMsg authSigMessage
			var _, err2 = ser.DecodeReaderWithType(sc, &_recvMsg, 1024*1024) // TODO
			if err2 != nil {
				return nil, err2, true // abort
			} else {
				return _recvMsg, nil, false
			}
		},
	)

	// If error:
	if trs.FirstError() != nil {
		err = trs.FirstError()
		return
	}

	var _recvMsg = trs.FirstValue().(authSigMessage)
	return _recvMsg, nil
}

//--------------------------------------------------------------------------------

// sha256
func hash32(input []byte) (res *[32]byte) {
	hasher := sha256.New()
	hasher.Write(input) // nolint: errcheck, gas
	resSlice := hasher.Sum(nil)
	res = new([32]byte)
	copy(res[:], resSlice)
	return
}

// We only fill in the first 20 bytes with ripemd160
func hash24(input []byte) (res *[24]byte) {
	hasher := ripemd160.New()
	hasher.Write(input) // nolint: errcheck, gas
	resSlice := hasher.Sum(nil)
	res = new([24]byte)
	copy(res[:], resSlice)
	return
}

// increment nonce big-endian by 2 with wraparound.
func incr2Nonce(nonce *[24]byte) {
	incrNonce(nonce)
	incrNonce(nonce)
}

// increment nonce big-endian by 1 with wraparound.
func incrNonce(nonce *[24]byte) {
	for i := 23; 0 <= i; i-- {
		nonce[i]++
		if nonce[i] != 0 {
			return
		}
	}
}
